package shardkv

import (
	"6.5840/labrpc"
	"6.5840/shardctrler"
	"bytes"
	"log"
	"sync/atomic"
	"time"
)
import "6.5840/raft"
import "sync"
import "6.5840/labgob"

const (
	WORKING int = 0
	MISSING int = 1
	ADDING  int = 2
)

const (
	UPDATECONFIG    = 100 * time.Millisecond
	SENDSHARDS      = 150 * time.Millisecond
	HandleOpTimeOut = time.Millisecond * 2000
)

const Debug = true

func DPrintf(format string, a ...interface{}) (n int, err error) {
	if Debug {
		log.Printf(format, a...)
	}
	return
}

type OpType string

const (
	OpGet    OpType = "Get"
	OpPut    OpType = "Put"
	OpAppend OpType = "Append"
)

type Op struct {
	// Your definitions here.
	// Field names must start with capital letters,
	// otherwise RPC will break.
	OpCmd    OpType
	Key      string
	Value    string
	Seq      uint64
	ClientId int64
}

type ShardKV struct {
	mu           sync.Mutex
	me           int
	rf           *raft.Raft
	dead         int32
	applyCh      chan raft.ApplyMsg
	make_end     func(string) *labrpc.ClientEnd
	gid          int
	ctrlers      []*labrpc.ClientEnd
	maxraftstate int // snapshot if log grows this big

	// Your definitions here.
	persister *raft.Persister
	db        map[string]string
	table     map[int64]uint64
	mck       *shardctrler.Clerk

	waiCh      map[int]*chan Result // 等待commit信息的RPC handler通道
	historyMap map[int64]*Result    // 记录clerk的最高序列号和对应的结构result

	lastConfig  shardctrler.Config
	config      shardctrler.Config
	shardsState map[int]int
	bytes       int
}

type Result struct {
	LastSeq uint64
	Err     Err
	Value   string
	ResTerm int // 记录commit时被apply的Term，可能与Start时相比发生变化，要返回给客户端
}

func (kv *ShardKV) checkShard(str string) (bool, bool) {
	shard := key2shard(str)
	if kv.config.Shards[shard] != kv.gid {
		return false, true
	}
	if kv.shardsState[shard] != WORKING {
		return true, false
	}
	return true, true
}

func (kv *ShardKV) HandleOp(opArgs *Op) (res Result) {
	kv.mu.Lock()

	//if history, ok := kv.historyMap[opArgs.ClientId]; ok && history.LastSeq == opArgs.Seq {
	//	// 当前操作已经执行过，直接从操作历史中获取结果
	//	kv.mu.Unlock()
	//	DPrintf("Leader %v handle operation: Client %v Seq %v Request %s(%s, %s) from history.\n", kv.me, opArgs.ClientId, opArgs.Seq, opArgs.OpCmd, opArgs.Key, opArgs.Value)
	//	return *history
	//}

	if shardMatch, shardReady := kv.checkShard(opArgs.Key); !shardMatch || !shardReady {
		if !shardMatch {
			res.Err = ErrWrongGroup
		} else if !shardReady {
			res.Err = ErrShardNotReady
		}
		kv.mu.Unlock()
		return
	}

	DPrintf("Leader %v handle operation: Client %v Seq %v Request %s(%s, %s) calls Start\n", kv.me, opArgs.ClientId, opArgs.Seq, opArgs.OpCmd, opArgs.Key, opArgs.Value)
	startIndex, startTerm, isLeader := kv.rf.Start(*opArgs)

	kv.mu.Unlock()

	if !isLeader {
		return Result{Err: ErrWrongLeader, Value: ""}
	}

	kv.mu.Lock()
	newCh := make(chan Result)
	kv.waiCh[startIndex] = &newCh
	kv.mu.Unlock()

	defer func() {
		kv.mu.Lock()
		delete(kv.waiCh, startIndex)
		close(newCh)
		kv.mu.Unlock()
	}()

	select {
	case <-time.After(HandleOpTimeOut): // 将处理时间超时设置为2s
		res.Err = ErrHandleOpTimeOut
		DPrintf("server %v Client %v Seq %v Timeout\n", kv.me, opArgs.ClientId, opArgs.Seq)
		return
	case msg, success := <-newCh:
		if success && msg.ResTerm == startTerm {
			res = msg
			res.Err = OK
			DPrintf("server %v HandleOp: Client %v Seq %v handle operation %s(%s, %s) successfully, res=%s", kv.me, opArgs.ClientId, opArgs.Seq, opArgs.OpCmd, opArgs.Key, opArgs.Value, res.Value)
			return
		} else if !success {
			DPrintf("server %v HandleOp: Client %v Seq %v closed chain, args.OpCmd=%v, args.Key=%s", kv.me, opArgs.ClientId, opArgs.Seq, opArgs.OpCmd, opArgs.Key)
			res.Err = ErrChanClosed
			return
		} else {
			DPrintf("server %v HandleOp: Client %v Seq %v Leader Term changed, res.ResTerm=%v startTerm=%v", kv.me, opArgs.ClientId, res.ResTerm, res.ResTerm, startTerm)
			res.Err = ErrLeaderOutdated
			return
		}
	}
}

func (kv *ShardKV) Get(args *GetArgs, reply *GetReply) {
	// Your code here.
	opArgs := &Op{
		OpCmd:    OpGet,
		Key:      args.Key,
		Seq:      args.Seq,
		ClientId: args.ClientId,
	}

	res := kv.HandleOp(opArgs)
	reply.Err = res.Err
	reply.Value = res.Value
}

func (kv *ShardKV) PutAppend(args *PutAppendArgs, reply *PutAppendReply) {
	// Your code here.
	opArgs := &Op{
		OpCmd:    OpPut,
		Key:      args.Key,
		Value:    args.Value,
		Seq:      args.Seq,
		ClientId: args.ClientId,
	}
	if args.Op == "Append" {
		opArgs.OpCmd = OpAppend
	}

	res := Result{}
	if opArgs.Seq <= kv.table[opArgs.ClientId] {
		res.Err = OK
		DPrintf("%s C%d seq %d aready done at S%d", opArgs.OpCmd, opArgs.ClientId, opArgs.Seq, kv.me)
		return
	}

	res = kv.HandleOp(opArgs)
	reply.Err = res.Err
}

func (kv *ShardKV) PushShard(args *PushShardArgs, reply *PushShardReply) {
	if kv.isLeader() == false {
		reply.Err = ErrWrongLeader
		return
	}
	kv.mu.Lock()
	if kv.config.Num < args.Num {
		reply.Err = OK
		kv.mu.Unlock()
		return
	}

	if kv.config.Num > args.Num {
		reply.Err = OK
		kv.mu.Unlock()
		go kv.sendDeleteShardRPC(args)
		return
	}
	if kv.config.Num == args.Num {
		pushed := false
		for _, shard := range args.Shards {
			if kv.shardsState[shard] == WORKING {
				pushed = true
			}
		}
		if pushed {
			reply.Err = OK
			kv.mu.Unlock()
			go kv.sendDeleteShardRPC(args)
			return
		}
	}
	kv.mu.Unlock()
	index, _, isLeader := kv.rf.Start(*args)

	if !isLeader {
		reply.Err = ErrWrongLeader
		return
	}
	ch := make(chan Result)
	kv.mu.Lock()
	kv.waiCh[index] = &ch
	kv.mu.Unlock()

	select {
	case res := <-ch:
		if res.Err == OK {
			reply.Err = res.Err
			go kv.sendDeleteShardRPC(args)
			return
		} else {
			reply.Err = res.Err
			return
		}
	case <-time.After(50 * time.Millisecond):
		reply.Err = ErrHandleOpTimeOut
		return
	}
}

func (kv *ShardKV) DeleteShard(args *DeleteShardArgs, reply *DeleteShardReply) {
	if kv.isLeader() == false {
		reply.Err = ErrWrongLeader
		return
	}
	kv.mu.Lock()
	if kv.config.Num > args.Num {
		reply.Err = OK
		kv.mu.Unlock()
		return
	}

	if kv.config.Num == args.Num {
		deleted := false
		for _, shard := range args.Shards {
			if kv.shardsState[shard] == WORKING {
				deleted = true
			}
		}
		if deleted {
			reply.Err = OK
			kv.mu.Unlock()
			return
		}
	}
	kv.mu.Unlock()
	index, _, isLeader := kv.rf.Start(*args)
	if !isLeader {
		reply.Err = ErrWrongLeader
		return
	}
	ch := make(chan Result)
	kv.mu.Lock()
	kv.waiCh[index] = &ch
	kv.mu.Unlock()

	select {
	case res := <-ch:
		reply.Err = res.Err
		return
	case <-time.After(50 * time.Millisecond):
		reply.Err = ErrHandleOpTimeOut
		return
	}
}

func (kv *ShardKV) isRepeated(clientId int64, seqNo uint64) bool {
	value, ok := kv.table[clientId]
	if ok && value >= seqNo {
		return true
	}
	return false
}

func (kv *ShardKV) isLeader() bool {
	_, isleader := kv.rf.GetState()
	return isleader
}

func (kv *ShardKV) readyForNewConfig() bool {
	for _, v := range kv.shardsState {
		if v != WORKING {
			return false
		}
	}
	return true
}

func (kv *ShardKV) updateConfig() {
	for kv.killed() == false {
		kv.mu.Lock()
		if !kv.readyForNewConfig() || !kv.isLeader() {
			kv.mu.Unlock()
			time.Sleep(UPDATECONFIG)
			continue
		}
		configNum := kv.config.Num
		kv.mu.Unlock()
		go func() {
			config := kv.mck.Query(configNum + 1)
			if config.Num > configNum {
				kv.rf.Start(config)
			}
		}()
		time.Sleep(UPDATECONFIG)

	}
}

func (kv *ShardKV) updateShardsState() {
	if kv.lastConfig.Num == 0 {
		return
	}

	for shard, gid := range kv.config.Shards {
		lastGid := kv.lastConfig.Shards[shard]
		if lastGid == kv.gid && gid != kv.gid {
			kv.shardsState[shard] = MISSING
		} else if lastGid != 0 && lastGid != kv.gid && gid == kv.gid {
			kv.shardsState[shard] = ADDING
		}
	}
}

func (kv *ShardKV) hasMissingShards() bool {
	for _, v := range kv.shardsState {
		if v == MISSING {
			return true
		}
	}
	return false
}

func (kv *ShardKV) preparePushShardArgs() []PushShardArgs {
	gs := make(map[int][]int)
	for shard, gid := range kv.config.Shards {
		lastGid := kv.lastConfig.Shards[shard]
		if lastGid == kv.gid && gid != kv.gid {
			gs[gid] = append(gs[gid], shard)
		}
	}

	res := make([]PushShardArgs, 0)
	for gid, shards := range gs {
		shardskv := make(map[string]string)
		hash := make(map[int]bool)
		for _, shard := range shards {
			hash[shard] = true
		}
		for k, v := range kv.db {
			if hash[key2shard(k)] == true {
				shardskv[k] = v
			}
		}
		args := PushShardArgs{
			Shards: shards,
			Num:    kv.config.Num,
			Data:   shardskv,
			Src:    kv.lastConfig.Groups[kv.gid],
			Table:  copyTable(kv.table),
			Dst:    kv.config.Groups[gid],
		}
		//args.Shards = shards
		//args.Num = kv.config.Num
		//args.Data = shardskv
		//args.Src = kv.lastConfig.Groups[kv.gid]
		//args.Table = copyTable(kv.table)
		//args.Dst = kv.config.Groups[gid]
		res = append(res, args)
	}
	return res
}

func (kv *ShardKV) sendPushShardRPC(args PushShardArgs) {
	for _, server := range args.Dst {
		srv := kv.make_end(server)
		var reply PushShardReply
		ok := srv.Call("ShardKV.PushShard", &args, &reply)
		if ok && (reply.Err == OK) {
			return
		}
	}

}

func (kv *ShardKV) sendDeleteShardRPC(args1 *PushShardArgs) {
	args := DeleteShardArgs{
		Shards: args1.Shards,
		Dst:    args1.Src,
		Num:    args1.Num,
		Keys:   make([]string, 0),
	}

	for k, _ := range args1.Data {
		args.Keys = append(args.Keys, k)
	}
	for _, server := range args.Dst {
		srv := kv.make_end(server)
		var reply DeleteShardReply
		ok := srv.Call("ShardKV.DeleteShard", &args, &reply)
		if ok && (reply.Err == OK) {
			return
		}
	}

}

func (kv *ShardKV) sendShards() {
	for kv.killed() == false {
		kv.mu.Lock()
		if !kv.hasMissingShards() || !kv.isLeader() {
			kv.mu.Unlock()
			time.Sleep(SENDSHARDS)
			continue
		}
		argsarr := kv.preparePushShardArgs()
		kv.mu.Unlock()
		for _, args := range argsarr {
			go kv.sendPushShardRPC(args)
		}
		time.Sleep(SENDSHARDS)
	}
}

// the tester calls Kill() when a ShardKV instance won't
// be needed again. you are not required to do anything
// in Kill(), but it might be convenient to (for example)
// turn off debug output from this instance.
func (kv *ShardKV) Kill() {
	kv.rf.Kill()
	// Your code here, if desired.
}

func (kv *ShardKV) killed() bool {
	z := atomic.LoadInt32(&kv.dead)
	return z == 1
}

func (kv *ShardKV) execute() {
	for kv.killed() == false {
		msg := <-kv.applyCh
		if msg.CommandValid {
			res := Result{}
			if op, ok := msg.Command.(Op); ok {
				res = kv.applyOp(&op)
			} else if cmdConfig, ok := msg.Command.(shardctrler.Config); ok {
				kv.applyConfig(cmdConfig)
			} else if args, ok := msg.Command.(PushShardArgs); ok {
				res = kv.applyPushShard(args)
			} else if args, ok := msg.Command.(DeleteShardArgs); ok {
				res = kv.applyDeleteShard(args)
			}

			kv.mu.Lock()
			ch, ok := kv.waiCh[msg.CommandIndex]
			kv.mu.Unlock()
			if ok {
				// 信息发送
				func() {
					defer func() {
						if recover() != nil {
							// 出现panic时说明通道关闭，进行日志打印
							DPrintf("")
						}
					}()
					res.ResTerm = msg.SnapshotTerm

					*ch <- res
				}()
			}

			kv.mu.Lock()
			if kv.maxraftstate != -1 && kv.persister.RaftStateSize() >= kv.maxraftstate/100*98 { // 如果到最大值快照，会发生log未截断的错误
				// 容量到98%时生成快照
				snapshot := kv.Snapshot()
				kv.rf.Snapshot(msg.CommandIndex, snapshot)
			}
			kv.mu.Unlock()
			//kv.replyMsg(msg, res)
			//kv.snapshot(msg)
		} else {
			kv.mu.Lock()
			snapshot := msg.Snapshot
			kv.readSnapshot(snapshot)
			kv.mu.Unlock()
		}
	}
}

func (kv *ShardKV) applyOp(op *Op) Result {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	res := Result{}
	//res.LastSeq = op.Seq
	if op.OpCmd == OpGet {
		clientId := op.ClientId
		flag1, flag2 := kv.checkShard(op.Key)
		if !flag1 || !flag2 {
			res.Err = ErrShardNotReady
			return res
		}
		if op.Seq > kv.table[clientId] {
			kv.table[clientId] = op.Seq
		}
		if kv.table[clientId] == op.Seq {
			val, ok := kv.db[op.Key]
			if ok {
				res.Err = OK
				res.Value = val
			} else {
				res.Err = ErrNoKey
				res.Value = ""
			}
			return res
		}

	} else if op.OpCmd == OpPut || op.OpCmd == OpAppend {
		key := op.Key
		value := op.Value
		clientId := op.ClientId
		flag1, flag2 := kv.checkShard(op.Key)
		if !flag1 || !flag2 {
			res.Err = ErrShardNotReady
			return res
		}
		if kv.isRepeated(op.ClientId, op.Seq) {
			res.Err = ErrRepeatOperation
			return res
		}
		if op.OpCmd == "Put" {
			kv.db[key] = value
		} else {
			kv.db[key] = kv.db[key] + value
		}
		DPrintf("记录操作ClientId %v op.Seq %v", clientId, op.Seq)
		kv.table[clientId] = op.Seq
		res.Err = OK
		return res
	}
	return res
}

func (kv *ShardKV) applyConfig(cfg shardctrler.Config) {
	kv.mu.Lock()
	defer kv.mu.Unlock()

	if cfg.Num == kv.config.Num+1 && kv.readyForNewConfig() {
		kv.lastConfig = kv.config
		kv.config = cfg
		kv.updateShardsState()
	}
}

func (kv *ShardKV) applyPushShard(args PushShardArgs) Result {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	res := Result{}
	//res.LastSeq = args.Seq
	if args.Num < kv.config.Num {
		res.Err = OK
		return res
	}

	if args.Num > kv.config.Num {
		res.Err = ErrShardNotReady
		return res
	}

	pushed := true
	for _, shard := range args.Shards {
		if kv.shardsState[shard] != WORKING {
			pushed = false
		}
	}
	if pushed {
		res.Err = OK
		return res
	}
	for _, shard := range args.Shards {
		kv.shardsState[shard] = WORKING
	}
	for k, v := range args.Data {
		kv.db[k] = v
	}
	for k, v := range args.Table {
		if v > kv.table[k] {
			kv.table[k] = v
		}
	}
	return res
}

func (kv *ShardKV) applyDeleteShard(args DeleteShardArgs) Result {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	res := Result{}
	if args.Num < kv.config.Num {
		res.Err = OK
		return res
	}

	if args.Num > kv.config.Num {
		res.Err = ErrShardNotReady
		return res
	}

	if kv.config.Num == args.Num {
		deleted := false
		for _, shard := range args.Shards {
			if kv.shardsState[shard] == WORKING {
				deleted = true
			}
		}
		if deleted {
			res.Err = OK
			return res
		}
	}

	for _, shard := range args.Shards {
		kv.shardsState[shard] = WORKING
	}
	for _, k := range args.Keys {
		delete(kv.db, k)
	}
	return res

}

//func (kv *ShardKV) replyMsg(msg raft.ApplyMsg, res Result) {
//kv.mu.Lock()
//ch, ok := kv.waiCh[msg.CommandIndex]
//kv.mu.Unlock()
//if ok {
//	if msg.IsLeader {
//
//		select {
//		case ch <- res:
//		default:
//		}
//	} else {
//		select {
//		case ch <- ErrNotApplied:
//		default:
//		}
//	}
//}
//}

//func (kv *ShardKV) snapshot(msg raft.ApplyMsg) {
//	kv.mu.Lock()
//	if kv.maxraftstate > 0 && kv.persister.RaftStateSize() > kv.maxraftstate {
//		snapshot := kv.Snapshot()
//		go func(i int) {
//			kv.rf.Snapshot(i, snapshot)
//		}(msg.CommandIndex)
//	}
//	kv.mu.Unlock()
//}

// TODO: SnapShot
func (kv *ShardKV) Snapshot() []byte {
	w := new(bytes.Buffer)
	e := labgob.NewEncoder(w)

	e.Encode(kv.db)
	e.Encode(kv.table)
	//e.Encode(kv.historyMap)
	e.Encode(kv.lastConfig)
	e.Encode(kv.config)
	e.Encode(kv.shardsState)

	serverSate := w.Bytes()
	return serverSate
}

func (kv *ShardKV) readSnapshot(snapshot []byte) {
	if len(snapshot) == 0 {
		DPrintf("server %v read snapshot failed: snapshot is null", kv.me)
		return
	}

	r := bytes.NewBuffer(snapshot)
	d := labgob.NewDecoder(r)

	tmpDB := make(map[string]string)
	//tmpHistoryMap := make(map[int64]*Result)
	tmpTable := make(map[int64]uint64)
	tmpLastConfig := shardctrler.Config{}
	tmpConfig := shardctrler.Config{}
	tmpShardsState := make(map[int]int)

	if d.Decode(&tmpDB) != nil ||
		d.Decode(&tmpTable) != nil ||
		d.Decode(&tmpLastConfig) != nil ||
		d.Decode(&tmpConfig) != nil ||
		d.Decode(&tmpShardsState) != nil {
		DPrintf("server %v read snapshot failed", kv.me)
	} else {
		kv.db = tmpDB
		kv.table = tmpTable
		kv.lastConfig = tmpLastConfig
		kv.config = tmpConfig
		kv.shardsState = tmpShardsState
		//kv.historyMap = tmpHistoryMap
		DPrintf("server %v read snapshot succeed", kv.me)
	}
}

// servers[] contains the ports of the servers in this group.
//
// me is the index of the current server in servers[].
//
// the k/v server should store snapshots through the underlying Raft
// implementation, which should call persister.SaveStateAndSnapshot() to
// atomically save the Raft state along with the snapshot.
//
// the k/v server should snapshot when Raft's saved state exceeds
// maxraftstate bytes, in order to allow Raft to garbage-collect its
// log. if maxraftstate is -1, you don't need to snapshot.
//
// gid is this group's GID, for interacting with the shardctrler.
//
// pass ctrlers[] to shardctrler.MakeClerk() so you can send
// RPCs to the shardctrler.
//
// make_end(servername) turns a server name from a
// Config.Groups[gid][i] into a labrpc.ClientEnd on which you can
// send RPCs. You'll need this to send RPCs to other groups.
//
// look at client.go for examples of how to use ctrlers[]
// and make_end() to send RPCs to the group owning a specific shard.
//
// StartServer() must return quickly, so it should start goroutines
// for any long-running work.
func StartServer(servers []*labrpc.ClientEnd, me int, persister *raft.Persister, maxraftstate int, gid int, ctrlers []*labrpc.ClientEnd, make_end func(string) *labrpc.ClientEnd) *ShardKV {
	// call labgob.Register on structures you want
	// Go's RPC library to marshall/unmarshall.
	labgob.Register(Op{})
	labgob.Register(PushShardArgs{})
	labgob.Register(DeleteShardArgs{})
	labgob.Register(shardctrler.Config{})

	kv := new(ShardKV)
	kv.me = me
	kv.maxraftstate = maxraftstate
	kv.make_end = make_end
	kv.gid = gid
	kv.ctrlers = ctrlers

	// Your initialization code here.
	kv.mck = shardctrler.MakeClerk(kv.ctrlers)

	// Use something like this to talk to the shardctrler:
	// kv.mck = shardctrler.MakeClerk(kv.ctrlers)

	kv.applyCh = make(chan raft.ApplyMsg)
	kv.rf = raft.Make(servers, me, persister, kv.applyCh)

	kv.table = make(map[int64]uint64)
	kv.db = make(map[string]string)
	kv.waiCh = make(map[int]*chan Result)
	kv.lastConfig = shardctrler.Config{}

	kv.config = shardctrler.Config{}
	kv.shardsState = make(map[int]int)
	for i := 0; i < 10; i++ {
		kv.shardsState[i] = 0
	}

	kv.persister = persister

	kv.mu.Lock()
	kv.readSnapshot(persister.ReadSnapshot())
	kv.mu.Unlock()

	go kv.execute()
	go kv.updateConfig()
	go kv.sendShards()

	return kv
}
